#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "itpack2d.h"

float laplace[9] = { 0.25, 0.5, 0.25,
		    0.5,   -3,  0.5,
		     0.25, 0.5, 0.25,};

float DD[9][9] = 
  { {  0,   0,   0,   0,   0,   0,   0,   0,   0,},
    {  0,  .5,   0,   0,  .5,   0,   0,   0,   0,},
    {  0,   0,   0,   0,   0,   0,   0,   0,   0,},
    {  0,   0,   0,  .5,  .5,   0,   0,   0,   0,},
    {  0, -.5,   0, -.5, -2., -.5,   0, -.5,   0,},
    {  0,   0,   0,   0,  .5,  .5,   0,   0,   0,},
    {  0,   0,   0,   0,   0,   0,   0,   0,   0,},
    {  0,   0,   0,   0,  .5,   0,   0,  .5,   0,},
    {  0,   0,   0,   0,   0,   0,   0,   0,   0,},
  };

float grad(float *data, int columns, int rows, int k,int l)
{
  // A = (sqrt(2)-1)/(2-sqrt(2))
  // B = 1/(2+4*A)
  // C = A*B
  float C = 0.14645, B = 0.20711;
  float x,y, r;
  x = ( -C*data[k*columns+l - columns -1] 
	-B*data[k*columns+l - columns]
	-C*data[k*columns+l - columns +1]
	+C*data[k*columns+l + columns -1] 
	+B*data[k*columns+l + columns]
	+C*data[k*columns+l + columns +1]);
  y = ( -C*data[k*columns+l -1 - columns] 
	-B*data[k*columns+l -1]
	-C*data[k*columns+l -1 + columns]
	+C*data[k*columns+l +1 - columns] 
	+B*data[k*columns+l +1]
	+C*data[k*columns+l +1 + columns]);
  r = sqrt(x*x+y*y);
  return r;
}


float lambda = 6.0;
float tau = 0.5;
float delta = 1.0;


float g(float s)
{
  s = s*s;
  if(s< 0.001) return 1.0;
  else return 1-expf(-3.315/powf(s/lambda,4));
}

#define JCOEF(i,j) (jcoef[(j)*(N)+(i)])
#define COEF(i,j) (coef[(j)*(N)+(i)])
#define GCOEF(i,j) (gcoef[(j)*(N)+(i)])
#define NOFFSET(k,l,x,y) (((k)+(x))*(columns)+(l)+(y))
#define JOFFSET(k,l,j) (((k)-1+(j)/3)*(columns)+(l)-1+(j)%3)

int CG(float *G, float * data, int columns, int rows)
{ 
  int k, l;
  //cauculate the G 
  for(k=1;k<rows-1;k++)
    for(l=1;l<columns-1;l++)
      G[k*columns+l] = g(grad(data, columns, rows, k,l));
  for(l=1;l<columns-1;l++)
    {
      G[l] = G[columns+l]; 
      G[rows*columns-l-1] = G[(rows-1)*columns-l-1];
    }
  for(k=1;k<rows-1;k++)
    {
      G[k*columns] = G[k*columns+1]; 
      G[(k+1)*columns-1] = G[(k+1)*columns-2];
    }
  G[0] = (G[1] + G[columns])/2;
  G[columns-1] = (G[columns-2]+ G[2*columns-1])/2;
  G[(rows-1)*columns] = (G[(rows-2)*columns] + G[(rows-1)*columns+1])/2;
  G[rows*columns-1] = (G[(rows-2)*columns-1] + G[rows*columns-2])/2;
  return 0;
}

int diffuse(float *data, int columns, int rows)
{
  int t,MAXT = 10;
  int N = columns*rows, NDIM=N, MAXNZ=9;
  int ITMAX = 20, NW=6*N+4*ITMAX, IER;
  float *coef, *gcoef, *WKSP, RPARM[12], *RHS, *G, *v;
  int *jcoef, *IWKSP, IPARM[12];
  int i,j,jj,k,l;

  jcoef = malloc(sizeof(int) * N * MAXNZ);
  coef = malloc(sizeof(float) * N * MAXNZ);
  gcoef = malloc(sizeof(float) * N * MAXNZ);
  RHS = malloc(sizeof(float) * N);
  v = malloc(sizeof(float) * N); //v
  G = malloc(sizeof(float) * N); //g(grad(v))

  WKSP = malloc(sizeof(float) * NW);
  IWKSP = malloc(sizeof(int) * N * 3);
  
  dfault_(IPARM, RPARM);
  IPARM[0] = ITMAX;
  IPARM[1] = 1; //output level
  IPARM[4] = 2; // nonsymmetic
  IPARM[11] = 1; //print DIGIT1 and DIGIT2
  
  for(i=0;i<columns;i++)
    {
      JCOEF(i,0) = i+1;
      JCOEF(N-i-1,0) = N-i;
      GCOEF(i, 0) = COEF(i, 0) =  1.;
      GCOEF(N-i-1,0) = COEF(N-i-1,0) = 1.;
      for(j=1;j<MAXNZ;j++)
	JCOEF(i,j) = 0, JCOEF(N-i-1,j) = 0;
      
    }
  for(i=1;i<rows-1;i++)
    {
      JCOEF(i*columns,0) = i*columns+1;
      GCOEF(i*columns,0) = COEF(i*columns,0) = 1.;
      JCOEF(i*columns+columns-1,0) = i*columns+columns;
      GCOEF(i*columns+columns-1,0) = COEF(i*columns+columns-1,0) = 1.;
      for(j=1;j<MAXNZ;j++)
	JCOEF(i*columns,j) = 0, JCOEF(i*columns+columns-1,j) = 0;
    }

  for(k=1;k<rows-1;k++)
    for(l=1;l<columns-1;l++)
      {
	for(j=0;j<MAXNZ;j++)
	  JCOEF(k*columns+l,j) = JOFFSET(k,l,j) + 1;
      }
  
  for(k=1;k<rows-1;k++)
    for(l=1;l<columns-1;l++)
      {
	for(j=0;j<MAXNZ;j++)
	  GCOEF(k*columns+l,j) = 0.;
	GCOEF(k*columns+l,4) = 1.;
	for(j=0;j<MAXNZ;j++)
	  GCOEF(k*columns+l,j) -= delta * laplace[j]; 
      }

  memcpy(v, data, sizeof(float)*N);
  
  // Note: SSORCG will sort the JCOEF and COEF in its own order.
  // So the calling order must desinged well
  // First JCOEF and GCOEF is in my own order, 
  // then run ssorcg, GCOEF's order changes.
  // Reset JCOEF and COEF to my own order, then run ssorcg.
  // Now JCOEF's order is SSORCG's, in the next iteration, 
  // GCOEF will not change, and JCOEF and GCOEF are both SSORCG's order
  // So the follow program can work.
  
  for(t = 0; t<MAXT; t++)
    {
      memcpy(RHS, data, sizeof(float)*N);


      ssorcg_(&N, &NDIM, &MAXNZ, jcoef, gcoef, RHS, v,
	      IWKSP, &NW, WKSP, IPARM, RPARM, &IER);
      /*
      for(i=0;i<N;i++)
	{
	  for(j=0;j<MAXNZ;j++)
	    printf("%4d", JCOEF(i,j));
	  printf("\n");
	}
      */
      
      for(k=1;k<rows-1;k++)
	for(l=1;l<columns-1;l++)
	  {
	    for(j=0;j<MAXNZ;j++)
	      JCOEF(k*columns+l,j) = JOFFSET(k,l,j) + 1;
	  }

      CG(G,v, columns, rows);

      for(k=1;k<rows-1;k++)
	for(l=1;l<columns-1;l++)
	  {
	    for(j=0;j<MAXNZ;j++)
	      COEF(k*columns+l,j) = 0.;
	    COEF(k*columns+l,4) = 1.;
	    for(j=0;j<MAXNZ;j++)
	      for(jj=0;jj<MAXNZ;jj++)
		COEF(k*columns+l,j) 
		  -= DD[j][jj]*G[JOFFSET(k,l,jj)]; 
	  }

      ssorcg_(&N, &NDIM, &MAXNZ, jcoef, coef, RHS, data,
	      IWKSP, &NW, WKSP, IPARM, RPARM, &IER);
    }

  free(jcoef);
  free(coef);
  free(gcoef);
  free(RHS);
  free(v);
  free(G);
  free(WKSP);
  free(IWKSP);
  printf("t: %d",t);
  return 0;
}
